{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sys\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "import os\n",
    "\n",
    "sys.path.append('../../aisuite')\n",
    "# Load from .env file if available\n",
    "load_dotenv(find_dotenv())\n",
    "os.environ['ALLOW_MULTI_TURN'] = 'true'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mock tool functions.\n",
    "def get_current_temperature(location: str, unit: str):\n",
    "    \"\"\"This is a short description of what the function does.\n",
    "\n",
    "    This is a longer description that can span\n",
    "    multiple lines and provide more details.\n",
    "\n",
    "    Args:\n",
    "        param1: Description of param1\n",
    "        param2: Description of param2\n",
    "    \"\"\"\n",
    "    return \"70\"\n",
    "\n",
    "def is_it_raining(location: str):\n",
    "    # Simulate fetching rain probability\n",
    "    return \"yes\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call the model with tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aisuite import Client\n",
    "\n",
    "client = Client()\n",
    "messages = [{\n",
    "    \"role\": \"user\",\n",
    "    \"content\": \"Can you plan a picnic for today afternoon in San Francisco? Check the temperature and if its raining.\"}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------- response from LLM ---------\n",
      "ChatCompletion(id='chatcmpl-AvuR3w6M83nWHL9sIO23pgvD0E5PF', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_PVhrHTUQ2qY7FDyPkCT54phU', function=Function(arguments='{\\n\"location\": \"San Francisco\",\\n\"unit\": \"Fahrenheit\"\\n}', name='get_current_temperature'), type='function')]))], created=1738364997, model='gpt-4-0613', object='chat.completion', service_tier='default', system_fingerprint=None, usage=CompletionUsage(completion_tokens=24, prompt_tokens=110, total_tokens=134, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))\n",
      "Executing tool:  get_current_temperature\n",
      "--------- tool_message to send to LLM ---------\n",
      "[{'role': 'tool', 'name': 'get_current_temperature', 'content': '\"70\"', 'tool_call_id': 'call_PVhrHTUQ2qY7FDyPkCT54phU'}]\n",
      "--------- response from LLM ---------\n",
      "ChatCompletion(id='chatcmpl-AvuR4hpiCgXFxRMNLg1Scv1UD2JlR', choices=[Choice(finish_reason='tool_calls', index=0, logprobs=None, message=ChatCompletionMessage(content=None, refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_Fvy3JUhIbVqzV0Nb0QfSnWQC', function=Function(arguments='{\\n\"location\": \"San Francisco\"\\n}', name='is_it_raining'), type='function')]))], created=1738364998, model='gpt-4-0613', object='chat.completion', service_tier='default', system_fingerprint=None, usage=CompletionUsage(completion_tokens=18, prompt_tokens=145, total_tokens=163, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))\n",
      "Executing tool:  is_it_raining\n",
      "--------- tool_message to send to LLM ---------\n",
      "[{'role': 'tool', 'name': 'is_it_raining', 'content': '\"yes\"', 'tool_call_id': 'call_Fvy3JUhIbVqzV0Nb0QfSnWQC'}]\n",
      "--------- response from LLM ---------\n",
      "ChatCompletion(id='chatcmpl-AvuR7puWSJLEpCNIZLLkF40gYJp3c', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content=\"I'm sorry, it seems like it will be raining this afternoon in San Francisco. You might want to plan your picnic for another day. Also the temperature is forecasted to be around 70 degrees Fahrenheit.\", refusal=None, role='assistant', audio=None, function_call=None, tool_calls=None))], created=1738365001, model='gpt-4-0613', object='chat.completion', service_tier='default', system_fingerprint=None, usage=CompletionUsage(completion_tokens=44, prompt_tokens=175, total_tokens=219, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))\n",
      "I'm sorry, it seems like it will be raining this afternoon in San Francisco. You might want to plan your picnic for another day. Also the temperature is forecasted to be around 70 degrees Fahrenheit.\n"
     ]
    }
   ],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=\"openai:gpt-4\", messages=messages, tools=[get_current_temperature, is_it_raining], max_turns=4)\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=\"anthropic:claude-3-5-sonnet-20241022\", messages=messages, tools=[get_current_temperature, is_it_raining], max_turns=4)\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ChatCompletionMessage(content=None, refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_PVhrHTUQ2qY7FDyPkCT54phU', function=Function(arguments='{\\n\"location\": \"San Francisco\",\\n\"unit\": \"Fahrenheit\"\\n}', name='get_current_temperature'), type='function')]),\n",
      " {'content': '\"70\"',\n",
      "  'name': 'get_current_temperature',\n",
      "  'role': 'tool',\n",
      "  'tool_call_id': 'call_PVhrHTUQ2qY7FDyPkCT54phU'},\n",
      " ChatCompletionMessage(content=None, refusal=None, role='assistant', audio=None, function_call=None, tool_calls=[ChatCompletionMessageToolCall(id='call_Fvy3JUhIbVqzV0Nb0QfSnWQC', function=Function(arguments='{\\n\"location\": \"San Francisco\"\\n}', name='is_it_raining'), type='function')]),\n",
      " {'content': '\"yes\"',\n",
      "  'name': 'is_it_raining',\n",
      "  'role': 'tool',\n",
      "  'tool_call_id': 'call_Fvy3JUhIbVqzV0Nb0QfSnWQC'},\n",
      " ChatCompletionMessage(content=\"I'm sorry, it seems like it will be raining this afternoon in San Francisco. You might want to plan your picnic for another day. Also the temperature is forecasted to be around 70 degrees Fahrenheit.\", refusal=None, role='assistant', audio=None, function_call=None, tool_calls=None)]\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint \n",
    "pprint(response.choices[0].intermediate_messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aisuite import Tools\n",
    "tools = Tools(tools=[get_current_temperature, is_it_raining])\n",
    "tools.tools()\n",
    "# tools.add_description(\"is_it_raining\", \"Use this function to understand if it is going to rain or not\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = append(messages, response.choices[0].intermediate_messages)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
